# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import getpass
import logging
import os
import yaml

from absl import (flags, app)

from coin.base.timestamp import get_timestamp
from coin.tool.slurm.job_group import JobGroup
from coin.tool.slurm.slurm import SlurmJobManager, SlurmJob

FLAGS = flags.FLAGS


def form_coin_repo_command(pyrunner_cmd):
  username = getpass.getuser()
  coin_run_path = f"/remote/iosg/home/{username}/workspace/coin/python"
  venv_path = ''
  cmd = ' && '.join([
      'env -i',
      f'cd {coin_run_path}',
      f'export TZ=UTC',
      f'source {venv_path}/bin/activate',
      f'./pyrunner {pyrunner_cmd}'
      'deactivate'
  ])
  return cmd


def get_noti_command(message):
  command = f"noti -S slurm -s @xguo {message}"
  return command


def get_slurm_working_dir(working_dir, prefix):
  timestamp = get_timestamp()
  return os.path.join(working_dir, f"{prefix}_{timestamp}")


def run_on_slurm(jobs_config):
  job_name = jobs_config['name']
  working_dir = jobs_config['working_dir']
  slurm_working_dir = get_slurm_working_dir(working_dir, job_name)
  logging.info('Slurm working dir: %s', slurm_working_dir)
  sim_jobs = []

  for single_job_config in jobs_config['jobs']:
    name = single_job_config['name']
    command = single_job_config['command']
    mem = single_job_config.get('mem')
    cpu = single_job_config.get('cpu')
    job = SlurmJob(name, command, mem=mem, cpu=cpu)
    sim_jobs.append(job)

  job_group = JobGroup(job_name, sim_jobs)
  trigger = SlurmJob('trigger', 'echo 0', trigger_job=True)
  job_group.start_trigger.set_afterok(trigger)
  job_manager = SlurmJobManager(
      slurm_working_dir,
      slurm_job_name_prefix='slurm',
      slurm_login_node='hive-29',
  )

  slack_job = SlurmJob('slack', get_noti_command('Done'))
  slack_job.set_afterok(job_group.end_trigger)
  all_jobs = job_group.get_all_jobs() + [trigger, slack_job]
  logging.info('Number of jobs: %s', len(all_jobs))
  job_manager.prepare_jobs(all_jobs)
  job_manager.start_job(trigger)


def main(_):
  assert FLAGS.config
  with open(FLAGS.config) as infile:
    config = yaml.load(infile)
  FLAGS.slurm_job_name_prefix = 'slurm'
  run_on_slurm(config)


if __name__ == '__main__':
  flags.DEFINE_string('config', None, 'Job configurations file.')
  app.run(main)
